import bagel
import datetime
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from sample.config.global_data import font_dict
import os
from pylab import *
import locale

locale.setlocale(locale.LC_CTYPE, 'chinese')
mpl.rcParams['font.sans-serif'] = ['SimHei']


def _expand(a: np.ndarray) -> np.ndarray:
    ret = np.copy(a)
    for i in range(length := len(a)):
        if a[i] == 1:
            if i - 1 >= 0:
                ret[i - 1] = 1
            if i + 1 < length:
                ret[i + 1] = 1
    return ret


def plot_res(kpi: bagel.data.KPI, change, bagel_data, test_file):
    control_test_file = test_file.replace('219', '220')
    control_test_kpi = None
    if os.path.exists(control_test_file):
        control_test_kpi = bagel.utils.load_kpi(control_test_file)
    fig = plt.figure(figsize=(80, 50))
    plt.xticks([])
    plt.yticks([])
    kpi_path = kpi.name.split(os.sep)
    plt.title(f'{kpi_path[-2]}_{kpi_path[-1]}', fontdict=font_dict)

    ax_value = fig.add_subplot(4, 1, 1)
    ax_value.set_ylabel('value', fontdict=font_dict)
    ax_value.set_xlabel('time', fontdict=font_dict)
    ax_value.tick_params(labelsize=50)

    ax_expectation = fig.add_subplot(4, 1, 2)
    ax_expectation.set_ylabel('expectation', fontdict=font_dict)
    ax_expectation.set_xlabel('time', fontdict=font_dict)
    ax_expectation.tick_params(labelsize=50)

    ax_score = fig.add_subplot(4, 1, 3)
    ax_score.set_ylabel('score', fontdict=font_dict)
    ax_score.set_xlabel('time', fontdict=font_dict)
    ax_score.tick_params(labelsize=50)

    ax_median = fig.add_subplot(4, 1, 4)
    ax_median.set_ylabel('median_diff', fontdict=font_dict)
    ax_median.set_xlabel('time', fontdict=font_dict)
    ax_median.tick_params(labelsize=50)

    x = [datetime.datetime.fromtimestamp(timestamp) for timestamp in kpi.timestamps]
    y_eval_anomaly = np.copy(kpi.raw_values)
    y_eval_anomaly[_expand(bagel_data['eval_label']) == 0] = np.inf
    y_real_anomaly = np.copy(kpi.raw_values)    
    y_real_anomaly[_expand(kpi.labels) == 0] = np.inf
    y_eval_anomaly_score = np.copy(bagel_data['anomaly_scores'])
    y_eval_anomaly_score[_expand(bagel_data['eval_label']) == 0] = np.inf
    # kpi原始值
    if change['start_time'] is not None:
        ax_value.vlines(
            [datetime.datetime.fromtimestamp(change['start_time']), datetime.datetime.fromtimestamp(change['end_time'])],
            np.min(kpi.raw_values), np.max(kpi.raw_values), colors='gold', linewidth=5)
    ax_value.plot(x, kpi.raw_values, color='blue', linewidth=4, label='kpi')
    ax_value.plot(x, y_real_anomaly, color='orange', linewidth=4, label='real label')
    if control_test_kpi is not None:
        ax_value.plot([datetime.datetime.fromtimestamp(timestamp) for timestamp in control_test_kpi.timestamps], control_test_kpi.raw_values, color='green', linewidth=4, label='control kpi')
    # 预测输出范围
    if change['start_time'] is not None:
        ax_expectation.vlines(
            [datetime.datetime.fromtimestamp(change['start_time']), datetime.datetime.fromtimestamp(change['end_time'])],
            np.min(kpi.raw_values), np.max(kpi.raw_values), colors='gold', linewidth=5)
    ax_expectation.plot(x, kpi.raw_values, color='black', linewidth=4, label='kpi')
    ax_expectation.plot(x, y_eval_anomaly, color='red', linewidth=4, label='eval label')
    ax_expectation.plot(x, bagel_data['x_expectation'], color='blue', linewidth=4, label='x_expectation')
    ax_expectation.plot(x, bagel_data['low_expectation'], color='green', linewidth=4, label='low_expectation')
    ax_expectation.plot(x, bagel_data['high_expectation'], color='purple', linewidth=4, label='high_expectation')
    if control_test_kpi is not None:
        ax_expectation.plot([datetime.datetime.fromtimestamp(timestamp) for timestamp in control_test_kpi.timestamps], control_test_kpi.raw_values, color='yellow', linewidth=4, label='control kpi')

    # 异常分数
    if change['start_time'] is not None:
        ax_score.vlines(
            [datetime.datetime.fromtimestamp(change['start_time']), datetime.datetime.fromtimestamp(change['end_time'])],
            np.min(bagel_data['anomaly_scores']), np.max(bagel_data['anomaly_scores']), colors='gold', linewidth=5)
    ax_score.plot(x, bagel_data['anomaly_scores'], color='palegreen', linewidth=4, label='score')
    ax_score.plot(x, y_eval_anomaly_score, color='red', linewidth=4, label='eval label')
    ax_score.plot(x, bagel_data['thresholds'], color='purple', linewidth=4, label='threshold')
    # 中位差
    if change['start_time'] is not None:
        ax_median.vlines(
            [datetime.datetime.fromtimestamp(change['start_time']), datetime.datetime.fromtimestamp(change['end_time'])],
            np.min(bagel_data['median_diff_list']), np.max(bagel_data['median_diff_list']), colors='gold', linewidth=5)
    ax_median.plot(x, bagel_data['median_diff_list'], color='g', linewidth=4, label='median_diff')

    ax_value.legend(prop=font_dict, loc=1)
    ax_expectation.legend(prop=font_dict, loc=1)
    ax_score.legend(prop=font_dict, loc=1)
    ax_median.legend(prop=font_dict, loc=1)
    return fig


def plot_res_paper(kpi: bagel.data.KPI, change, bagel_data, test_file):
    control_test_file = test_file.replace('219', '220')
    control_test_kpi = None
    if os.path.exists(control_test_file):
        control_test_kpi = bagel.utils.load_kpi(control_test_file)
    fig = plt.figure(figsize=(100, 35))
    plt.xticks([])
    plt.yticks([])
    kpi_path = kpi.name.split(os.sep)
    plt.title(f'{kpi_path[-2]}_{kpi_path[-1]}'.replace('.csv', ''), fontdict=font_dict)

    ax_value = fig.add_subplot(2, 1, 1)
    ax_value.set_ylabel('value', fontdict=font_dict)
    ax_value.set_xlabel('time', fontdict=font_dict)
    ax_value.tick_params(labelsize=110)
    ax_value.xaxis.set_major_formatter(mdates.DateFormatter("%m月%d日%H时"))
    # ax_value.xaxis.set_major_locator(mdates.HourLocator(interval=6))
    ax_score = fig.add_subplot(2, 1, 2)
    ax_score.set_ylabel('score', fontdict=font_dict)
    ax_score.set_xlabel('time', fontdict=font_dict)
    ax_score.tick_params(labelsize=110)
    ax_score.xaxis.set_major_formatter(mdates.DateFormatter("%m月%d日%H时"))
    # ax_score.xaxis.set_major_locator(mdates.HourLocator(interval=6))
    x = [datetime.datetime.fromtimestamp(timestamp) for timestamp in kpi.timestamps]
    y_eval_anomaly = np.copy(kpi.raw_values)
    y_eval_anomaly[_expand(bagel_data['eval_label']) == 0] = np.inf
    y_real_anomaly = np.copy(kpi.raw_values)    
    y_real_anomaly[_expand(kpi.labels) == 0] = np.inf
    y_eval_anomaly_score = np.copy(bagel_data['anomaly_scores'])
    y_eval_anomaly_score[_expand(bagel_data['eval_label']) == 0] = np.inf
    # kpi原始值
    if change['start_time'] is not None:
        vertical_time_list = change['start_time']
        vertical_time_list.extend(change['end_time'])
        # ax_value.vlines(
        #     [datetime.datetime.fromtimestamp(ts) for ts in vertical_time_list],
        #     np.min(kpi.raw_values), np.max(kpi.raw_values), colors='black', linewidth=10)
    ax_value.plot(x, kpi.raw_values, color='blue', linewidth=15, label='study KPI')
    ax_value.plot(x, y_eval_anomaly, color='red', linewidth=15, label='anomaly label')
    if control_test_kpi is not None:
        ax_value.plot([datetime.datetime.fromtimestamp(timestamp) for timestamp in control_test_kpi.timestamps], control_test_kpi.raw_values, color='green', linewidth=8, label='control KPI')
    # 异常分数
    if change['start_time'] is not None:
        vertical_time_list = change['start_time']
        vertical_time_list.extend(change['end_time'])
        # ax_value.vlines(
        #     [datetime.datetime.fromtimestamp(ts) for ts in vertical_time_list],
        #     np.min(kpi.raw_values), np.max(kpi.raw_values), colors='black', linewidth=10)
    ax_score.plot(x, bagel_data['anomaly_scores'], color='green', linewidth=15, label='anomaly score')
    ax_score.plot(x, y_eval_anomaly_score, color='red', linewidth=15, label='anomaly label')
    ax_score.plot(x, bagel_data['thresholds'], color='purple', linewidth=8, label='score threshold')

    ax_value.legend(prop=font_dict, loc=1)
    ax_score.legend(prop=font_dict, loc=1)
    return fig